use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
use tokio::time;
use serde::{Serialize, Deserialize};
use anyhow::{Result, Context};
use log::{info, warn, error};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ServiceStatus {
    Healthy,
    Degraded,
    Unhealthy,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HealthStatus {
    pub service: String,
    pub status: ServiceStatus,
    pub last_check: String,
    pub message: String,
    pub recovery_attempts: u32,
}

pub struct HealthCheck {
    pub name: String,
    pub check_fn: Box<dyn Fn() -> Result<()> + Send + Sync>,
    pub recovery_fn: Option<Box<dyn Fn() -> Result<()> + Send + Sync>>,
    pub max_recovery_attempts: u32,
    pub check_interval: Duration,
}

pub struct HealthMonitor {
    checks: Vec<HealthCheck>,
    statuses: Arc<RwLock<Vec<HealthStatus>>>,
}

impl HealthMonitor {
    pub fn new() -> Self {
        Self {
            checks: Vec::new(),
            statuses: Arc::new(RwLock::new(Vec::new())),
        }
    }

    pub fn register_check(&mut self, check: HealthCheck) {
        let status = HealthStatus {
            service: check.name.clone(),
            status: ServiceStatus::Healthy,
            last_check: chrono::Utc::now().to_rfc3339(),
            message: "Service initialized".to_string(),
            recovery_attempts: 0,
        };
        
        let statuses = self.statuses.clone();
        tokio::spawn(async move {
            let mut statuses = statuses.write().await;
            statuses.push(status);
        });
        
        self.checks.push(check);
    }

    pub async fn start_monitoring(&self) {
        for (index, check) in self.checks.iter().enumerate() {
            let check_fn = check.check_fn.clone();
            let recovery_fn = check.recovery_fn.clone();
            let max_recovery_attempts = check.max_recovery_attempts;
            let check_interval = check.check_interval;
            let statuses = self.statuses.clone();
            let service_name = check.name.clone();
            
            tokio::spawn(async move {
                let mut interval = time::interval(check_interval);
                
                loop {
                    interval.tick().await;
                    
                    let result = (check_fn)();
                    let mut statuses = statuses.write().await;
                    
                    if let Some(status) = statuses.get_mut(index) {
                        status.last_check = chrono::Utc::now().to_rfc3339();
                        
                        match result {
                            Ok(_) => {
                                if status.status != ServiceStatus::Healthy {
                                    info!("Service {} recovered", service_name);
                                }
                                status.status = ServiceStatus::Healthy;
                                status.message = "Service is healthy".to_string();
                                status.recovery_attempts = 0;
                            },
                            Err(e) => {
                                warn!("Health check failed for {}: {}", service_name, e);
                                
                                if status.status == ServiceStatus::Healthy {
                                    status.status = ServiceStatus::Degraded;
                                    status.message = format!("Service degraded: {}", e);
                                } else if status.recovery_attempts < max_recovery_attempts {
                                    if let Some(ref recovery) = recovery_fn {
                                        info!("Attempting recovery for {}", service_name);
                                        status.recovery_attempts += 1;
                                        
                                        match recovery() {
                                            Ok(_) => {
                                                info!("Recovery attempt successful for {}", service_name);
                                                status.message = format!(
                                                    "Recovery attempt {} of {} successful", 
                                                    status.recovery_attempts, 
                                                    max_recovery_attempts
                                                );
                                            },
                                            Err(e) => {
                                                error!("Recovery attempt failed for {}: {}", service_name, e);
                                                status.message = format!(
                                                    "Recovery attempt {} of {} failed: {}", 
                                                    status.recovery_attempts, 
                                                    max_recovery_attempts, 
                                                    e
                                                );
                                                
                                                if status.recovery_attempts >= max_recovery_attempts {
                                                    status.status = ServiceStatus::Unhealthy;
                                                    status.message = format!(
                                                        "Service unhealthy after {} recovery attempts: {}", 
                                                        max_recovery_attempts, 
                                                        e
                                                    );
                                                }
                                            }
                                        }
                                    } else {
                                        status.status = ServiceStatus::Unhealthy;
                                        status.message = format!("Service unhealthy: {}", e);
                                    }
                                }
                            }
                        }
                    }
                }
            });
        }
    }

    pub async fn get_health_status(&self) -> Vec<HealthStatus> {
        self.statuses.read().await.clone()
    }

    pub async fn is_system_healthy(&self) -> bool {
        let statuses = self.statuses.read().await;
        !statuses.iter().any(|s| s.status == ServiceStatus::Unhealthy)
    }
}